#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <tss2/tss2_rc.h>
#include <tss2/tss2_sys.h>
#include <tss2/tss2_tcti.h>
#include <tss2/tss2_tcti_mssim.h>
#include <tss2/tss2_tctildr.h>
#include <tss2/tss2_tpm2_types.h>
#include <unistd.h>

#ifndef TCTI_NAME
#define TCTI_NAME "mssim"
#endif

TSS2_TCTI_CONTEXT* tcti_context = NULL;
TSS2_SYS_CONTEXT* sys_context = NULL;

// Record handles for cleanup
typedef struct HandleRecorder {
    TPM2_HANDLE handles[128];
    uint16_t count;
} HandleRecorder;


static HandleRecorder handle_recorder = {
    .count = 0, .handles = {TPM2_RH_NULL} // Initialize with NULL handle
};

// Check if the return code indicates success, clean up if not
#define CHECK_PASSED(rc)                                                                           \
    {                                                                                              \
        if ((rc) != TPM2_RC_SUCCESS) {                                                             \
            printf("FAILED! [%s:%u] %s\n", __FUNCTION__, __LINE__, Tss2_RC_Decode(rc));            \
            cleanup(rc);                                                                           \
        } else {                                                                                   \
            printf("PASSED! [%s:%u]\n", __FUNCTION__, __LINE__);                                   \
        }                                                                                          \
    }

// Initialize TCTI Context
TSS2_RC init_tcti_context(TSS2_TCTI_CONTEXT** tcti_context) {
    TSS2_RC rc = Tss2_TctiLdr_Initialize(TCTI_NAME, tcti_context);
    return rc;
}

// Initialize SAPI Context
TSS2_RC init_sys_context(TSS2_TCTI_CONTEXT* tcti_context, TSS2_SYS_CONTEXT** sys_context) {
    if (tcti_context == NULL) {
        return TSS2_SYS_RC_BAD_REFERENCE;
    }

    size_t sys_context_size = Tss2_Sys_GetContextSize(0);
    *sys_context = (TSS2_SYS_CONTEXT*)malloc(sys_context_size);
    if (*sys_context == NULL) {
        return TSS2_BASE_RC_MEMORY;
    }

    TSS2_ABI_VERSION abi_version = TSS2_ABI_VERSION_CURRENT;
    TSS2_RC rc = Tss2_Sys_Initialize(*sys_context, sys_context_size, tcti_context, &abi_version);
    if (rc != TSS2_RC_SUCCESS) {
        free(*sys_context);
        *sys_context = NULL;
    }

    return rc;
}

// Startup TPM
TSS2_RC startup_tpm(TSS2_SYS_CONTEXT* sys_context) {
    if (sys_context == NULL) {
        return TSS2_SYS_RC_BAD_REFERENCE;
    }

    TSS2_RC rc = Tss2_Sys_GetCapability(sys_context, NULL, TPM2_CAP_TPM_PROPERTIES,
                                        TPM2_PT_MANUFACTURER, 1, NULL, NULL, NULL);
    if (rc == TPM2_RC_INITIALIZE) {
        rc = Tss2_Sys_Startup(sys_context, TPM2_SU_CLEAR);
    }

    return rc;
}

// Cleanup resources
static void cleanup(int rc) {
    for (uint16_t i = 0; i < handle_recorder.count; i++) {
        if (handle_recorder.handles[i] != TPM2_RH_NULL) {
            TSS2_RC rc_flush = Tss2_Sys_FlushContext(sys_context, handle_recorder.handles[i]);
            if (rc_flush != TSS2_RC_SUCCESS) {
                printf("FlushContext Failed for handle 0x%08x\n", handle_recorder.handles[i]);
            }
        }
    }
    handle_recorder.count = 0;
    if (tcti_context != NULL) {
        Tss2_TctiLdr_Finalize(&tcti_context);
        free(tcti_context);
    }
    if (sys_context != NULL) {
        Tss2_Sys_Finalize(sys_context);
        free(sys_context);
    }
    exit(rc != TSS2_RC_SUCCESS);
    return;
}

// Create a symmetric key using AES in CFB mode
TSS2_RC create_key_aes_128_cfb(TSS2_SYS_CONTEXT* sys_context, const char* password,
                               TPM2_HANDLE* key_handle) {
    // CreatePrimary Input Parameters
    TSS2L_SYS_AUTH_COMMAND cmd_auths_array = {.count = 1, .auths = {{.sessionHandle = TPM2_RH_PW}}};
    TPM2B_PUBLIC in_public = {
        .size = 0,
        .publicArea =
            {
                .type = TPM2_ALG_SYMCIPHER,
                .nameAlg = TPM2_ALG_SHA256,
                .objectAttributes = TPMA_OBJECT_DECRYPT | TPMA_OBJECT_SIGN_ENCRYPT |
                                    TPMA_OBJECT_USERWITHAUTH | TPMA_OBJECT_SENSITIVEDATAORIGIN,
                .parameters.symDetail.sym =
                    {
                        .algorithm = TPM2_ALG_AES,
                        .keyBits.aes = 128,
                        .mode.sym = TPM2_ALG_CFB,
                    },
            },
    };
    UINT16 pwd_size = strlen(password);
    TPM2B_SENSITIVE_CREATE in_sensitive = {
        .sensitive.userAuth =
            {
                .size = pwd_size,
                .buffer = {},
            },
    };
    memcpy(in_sensitive.sensitive.userAuth.buffer, password, pwd_size);
    TPM2B_DATA outside_info = {
        .size = 0,
        .buffer = {},
    };
    TPML_PCR_SELECTION creation_pcr = {
        .count = 0,
    };

    // CreatePrimary Output Parameters
    TPM2B_PUBLIC out_public = {0};
    TPM2B_CREATION_DATA creation_data = {0};
    TPM2B_DIGEST creation_hash = {.size = sizeof(creation_hash.buffer)};
    TPMT_TK_CREATION creation_ticket = {0};
    TPM2B_NAME name = {0};
    TSS2L_SYS_AUTH_RESPONSE rsp_auths_array = {0};

    TSS2_RC rc = Tss2_Sys_CreatePrimary(sys_context, TPM2_RH_OWNER, &cmd_auths_array, &in_sensitive,
                                        &in_public, &outside_info, &creation_pcr, key_handle,
                                        &out_public, &creation_data, &creation_hash,
                                        &creation_ticket, &name, &rsp_auths_array);
    return rc;
};

// Encrypt data using the symmetric key
TSS2_RC encrypt(TSS2_SYS_CONTEXT* sys_context, const char* password, TPM2_HANDLE key_handle,
                TPM2B_MAX_BUFFER* in_data, TPM2B_MAX_BUFFER* out_data) {
    // EncryptDecrypt Input Parameters
    TSS2L_SYS_AUTH_COMMAND cmd_auths_array = {.count = 1,
                                              .auths[0] = {
                                                  .sessionHandle = TPM2_RH_PW,
                                                  .hmac = {.size = strlen(password)},
                                              }};
    memcpy(cmd_auths_array.auths[0].hmac.buffer, password, strlen(password));

    TPMI_ALG_CIPHER_MODE mode = TPM2_ALG_NULL;
    TPM2B_IV in_iv = {.size = 16};

    // EncryptDecrypt Output Parameters
    TPM2B_IV out_iv = {0};
    TSS2L_SYS_AUTH_RESPONSE rsp_auths_array = {0};

    TSS2_RC rc =
        Tss2_Sys_EncryptDecrypt2(sys_context, key_handle, &cmd_auths_array, in_data, TPM2_NO, mode,
                                 &in_iv, out_data, &out_iv, &rsp_auths_array);
    return rc;
}

// Decrypt data using the symmetric key
TSS2_RC decrypt(TSS2_SYS_CONTEXT* sys_context, const char* password, TPM2_HANDLE key_handle,
                TPM2B_MAX_BUFFER* in_data, TPM2B_MAX_BUFFER* out_data) {
    // EncryptDecrypt Input Parameters
    TSS2L_SYS_AUTH_COMMAND cmd_auths_array = {.count = 1,
                                              .auths[0] = {
                                                  .sessionHandle = TPM2_RH_PW,
                                                  .hmac = {.size = strlen(password)},
                                              }};
    memcpy(cmd_auths_array.auths[0].hmac.buffer, password, strlen(password));

    TPMI_ALG_CIPHER_MODE mode = TPM2_ALG_NULL;
    TPM2B_IV in_iv = {.size = 16};

    // EncryptDecrypt Output Parameters
    TPM2B_IV out_iv = {0};
    TSS2L_SYS_AUTH_RESPONSE rsp_auths_array = {0};

    TSS2_RC rc =
        Tss2_Sys_EncryptDecrypt2(sys_context, key_handle, &cmd_auths_array, in_data, TPM2_YES, mode,
                                 &in_iv, out_data, &out_iv, &rsp_auths_array);
    return rc;
}

int main(int argc, char* argv[]) {

    TSS2_RC rc;

    // Initialize TCTI contexts
    rc = init_tcti_context(&tcti_context);
    CHECK_PASSED(rc);

    // Initialize SAPI context
    rc = init_sys_context(tcti_context, &sys_context);
    CHECK_PASSED(rc);

    // Startup TPM
    rc = startup_tpm(sys_context);
    CHECK_PASSED(rc);

    // Create a symmetric key
    const char* password = "HelloWorld";
    TPM2_HANDLE key_handle;
    rc = create_key_aes_128_cfb(sys_context, password, &key_handle);
    CHECK_PASSED(rc)
    handle_recorder.handles[handle_recorder.count++] = key_handle; // Record the key handle

    // Encrypt data
    const char* plaintext = "This is a test message.";
    printf("Plain Text: %s\n", plaintext);
    TPM2B_MAX_BUFFER encrypt_in_data = {
        .size = strlen(plaintext),
        .buffer = {0},
    };
    memcpy(encrypt_in_data.buffer, plaintext, encrypt_in_data.size);
    TPM2B_MAX_BUFFER encrypt_out_data = {0};
    rc = encrypt(sys_context, password, key_handle, &encrypt_in_data, &encrypt_out_data);
    CHECK_PASSED(rc);

    // Decrypt data
    TPM2B_MAX_BUFFER decrypt_out_data = {0};
    rc = decrypt(sys_context, password, key_handle, &encrypt_out_data, &decrypt_out_data);
    CHECK_PASSED(rc);
    printf("Decrypted Text: %.*s\n", decrypt_out_data.size, decrypt_out_data.buffer); // Print decrypted datas

    // Cleanup resources
    cleanup(rc);

    return 0;
}
